from manimlib import *
import numpy as np


# 弹簧圈数
turns = 20

def Spring2D(start, end, axis):
    """
    创建弹簧，包括首尾节点（球体）和长度标签
    """
    start = np.array(start)
    end = np.array(end)

    global_length = np.linalg.norm(end - start)

    start = axis.coords_to_point(start[0], start[1])
    end = axis.coords_to_point(end[0], end[1])

    # 计算起始位置和终点位置之间的向量
    vector = end - start
    # 计算向量的长度
    length = np.linalg.norm(vector)
    # 计算向量的方向
    direction = vector / length

    # 横截面半径
    radius = global_length *0.05

    curve = mobject.functions.FunctionGraph(
        lambda t: axis.coords_to_point(t, radius * np.cos(2 * np.pi * turns * t), 0),  # x坐标  # y坐标  # z坐标
        t_range=[0, 1, 0.001],  # 绘制步长0.001
        color=YELLOW,
    )
    curve.set_style(stroke_width=1)

    # 将单元弹簧映射到真实坐标上
    curve.shift(start)  # 将弹簧移动到起始位置
    curve.stretch(global_length, 0, about_point=start)  # 根据向量的长度来缩放弹簧
    curve.rotate(np.arccos(direction[0]), about_point=start)  # 根据向量的方向来旋转弹簧

    # 在弹簧的两端添加两个球体
    circle_start = Circle(radius=radius * 0.2, color=RED).move_to(start)
    circle_start.set_fill(RED)
    circle_end = Circle(radius=radius * 0.2, color=RED).move_to(end)
    circle_end.set_fill(RED)

    # 创建弹簧长度的标签
    length_label = Text(str(global_length)).next_to(curve, UP)

    # 将弹簧和球体组合成一个整体
    return Group(curve, circle_start, circle_end, length_label)


def Spring3D(start, end, axis):
    """
    创建弹簧，包括首尾节点（球体）和长度标签
    """

    start = np.array(start)
    end = np.array(end)

    global_length = np.linalg.norm(end - start)

    start = axis.coords_to_point(start[0], start[1], start[2])
    end = axis.coords_to_point(end[0], end[1], end[2])

    # 计算起始位置和终点位置之间的向量
    vector = end - start
    # 计算向量的长度
    length = np.linalg.norm(vector)
    # 计算向量的方向
    direction = vector / length

    # 横截面半径
    radius = global_length *0.05

    curve = mobject.functions.ParametricCurve(
        lambda t: axis.coords_to_point(
            t, radius * np.cos(2 * np.pi * turns * t), radius * np.sin(2 * np.pi * turns * t)  # x坐标  # y坐标  # z坐标
        ),
        t_range=[0, 1, 0.005],  # 绘制步长0.001
        color=YELLOW
    )
    curve.set_style(stroke_width=10)

    # 将单元弹簧映射到真实坐标上
    curve.shift(start)  # 将弹簧移动到起始位置
    curve.stretch(global_length, 0, about_point=start)  # 根据向量的长度来缩放弹簧
    curve.rotate(np.arccos(direction[0]), about_point=start)  # 根据向量的方向来旋转弹簧

    # 在弹簧的两端添加两个球体
    ball_start = Sphere(radius=radius * 0.2, color=RED).move_to(start)
    ball_end = Sphere(radius=radius * 0.2, color=RED).move_to(end)

    # 将弹簧和球体组合成一个整体
    group = Group(curve, ball_start, ball_end)
    group.set_width(global_length)

    return group


def Bar3D(start, end, axis):
    """
    创建杆
    """
    
    start = np.array(start)
    end = np.array(end)

    global_length = np.linalg.norm(end - start)

    start = axis.coords_to_point(start[0], start[1], start[2])
    end = axis.coords_to_point(end[0], end[1], end[2])

    # 计算起始位置和终点位置之间的向量
    vector = end - start
    # 计算向量的长度
    length = np.linalg.norm(vector)
    # 计算向量的方向
    direction = vector / length

    # 横截面半径
    radius = global_length *0.05

    cylinder = Group(
        Cylinder(radius=radius, height=length, color=YELLOW),
        Circle(radius=radius).move_to(start).set_fill(YELLOW, opacity=1).set_style(stroke_width=0),
        Circle(radius=radius).shift(end).set_fill(YELLOW, opacity=1).set_style(stroke_width=0),
    )

    # 将单元弹簧映射到真实坐标上
    cylinder.shift(start)  # 将弹簧移动到起始位置
    cylinder.stretch(global_length, 0, about_point=start)  # 根据向量的长度来缩放弹簧
    cylinder.rotate(np.arccos(direction[0]), about_point=start)  # 根据向量的方向来旋转弹簧

    # 在弹簧的两端添加两个球体
    ball_start = Sphere(radius=radius * 0.2, color=RED).move_to(start)
    ball_end = Sphere(radius=radius * 0.2, color=RED).move_to(end)

    # 将弹簧和球体组合成一个整体
    group = Group(cylinder, ball_start, ball_end)
    group.set_width(global_length)

    return group